#!/usr/bin/python
import sys
import random

if __name__ == '__main__':
    if len(sys.argv) < 2:
	print 'usage: <fold> [nfold=5]'
	exit(-1)
    if len( sys.argv ) > 2:
        nfold = int( sys.argv[2] )
    else:
        nfold = 5
    fold = int( sys.argv[1] )
    assert fold > 0 and fold  <= nfold  

    random.seed( 0 )
    mints = min( int(l.split('::')[3]) for l in open( '../ml-1m/ratings.dat') )
    print 'fold=%d/%d' % ( fold, nfold )

    maxuid = 0
    maxts  = 0
    maxiid = 0
    ftr = open( 'train.txt', 'w' )
    fte = open( 'test.txt', 'w' )
    fi  = open( '../ml-1m/ratings.dat', 'r' ) 
    for line in fi:
        arr = line.split('::')
        uid = int( arr[0] ) - 1
        iid = int( arr[1] ) - 1
        score = int( arr[2] ) - 1 
        ts    = int( arr[3] ) - mints
        maxts  = max( ts, maxts )
        maxuid = max( uid, maxuid )
        maxiid = max( iid, maxiid )
        if random.randint( 1, nfold ) == fold:            
            fte.write( '%d\t%d\t%d\t%d\n' % ( uid, iid, score, ts ) )   
        else: 
            ftr.write( '%d\t%d\t%d\t%d\n' % ( uid, iid, score, ts ) )          
    ftr.close()
    fte.close()
    

    print 'maxuid=%d, maxiid=%d, maxts=%d' % ( maxuid+1, maxiid+1, maxts + 1 )
    sc = [ float(l.split('::')[2]) - 1  for l in open( '../ml-1m/ratings.dat')]      
    print 'score=[0-4], mean = %f' % ( sum(sc) / len(sc) )

